package graphx

import algorithm.ShortestPathAlgorithm
import org.apache.log4j.{Level, Logger}
import org.apache.spark.graphx.{Edge, Graph, VertexId}
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.types._

import scala.reflect.io.Directory

object shortestPathExample extends App {
  //定义节点的属性
  case class VertexAttr(name: String, latitude: Double, longitude: Double, population: Int)

  //设置日志级别
  Logger.getLogger("org").setLevel(Level.ERROR)
  var resourcesPath = "file:///" + Directory.Current.get.toString()

  val spark = SparkSession.builder()
    .appName("shortest-path-scala")
    .master("local[2]")
    .getOrCreate()
  //获取节点
  val gNodes: RDD[(VertexId, VertexAttr)] = spark.read.format("csv")
    .schema(StructType(Array(
      StructField("id", IntegerType, nullable = true),
      StructField("name", StringType, nullable = true),
      StructField("latitude", DoubleType, nullable = true),
      StructField("longitude", DoubleType, nullable = true),
      StructField("population", IntegerType, nullable = true))))
    .option("header", "true")
    .load(resourcesPath + "/data/transport-nodes.csv")
    .rdd
    .map(row => {
      (row.getInt(0),
        VertexAttr(row.getAs[String]("name"),
          row.getAs[Double]("latitude"),
          row.getAs[Double]("longitude"),
          row.getAs[Integer]("population")))
    }) //转换成rdd

  //获取关系 (这里简化了指定schema)
  val gRels = spark.read.format("csv")
    .schema(StructType(Array(
      StructField("src", IntegerType, nullable = true),
      StructField("dst", IntegerType, nullable = true),
      StructField("relationship", StringType, nullable = true),
      StructField("cost", DoubleType, nullable = true))))
    .option("header", "true")
    .load(resourcesPath + "/data/transport-relationships.csv")

  //合成完整的关系
  val gRelationships: RDD[Edge[Double]] = gRels
    .rdd
    .map(row => {
      Edge(row.getInt(0), row.getInt(1), row.getAs[Double]("cost"))
    })

  var graph = Graph(gNodes, gRelationships)

  ShortestPathAlgorithm.dijkstraWithTrace(graph,3L,direct = false).vertices.map(_._2).collect().foreach(data
  =>{
    println(data._1,data._2,data._3)
  })

  ShortestPathAlgorithm.dijkstraWithTrace(graph,3L,direct = true).vertices.map(_._2).collect().foreach(data
  =>{
    println(data._1,data._2,data._3)
  })


  ShortestPathAlgorithm.pregel(graph,3L,direct = false,20).vertices.map(_._2).collect().foreach(data
  =>{
    println(data._1,data._2)
  })
}
